#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <netinet/in.h>
#include <memory.h>
#include "list_dir.h"
#include "RetrievePcap.h"
#include "mysql/mysql.h"

MYSQL my_connection;

//construct a new flow with the new packet not belong to any existent flow
void process_flow(FILE* PcapLog,struct packet *newPacket,char Class[50],struct statistic Statistic)
{
    char query[300];
    if(strcmp(Class,"\"UNKNOWN\""))
    {
        sprintf(query,"call AddValue(%s,%ld,%d,%ld,%d,%ld,%d,%d,%s)",Statistic.filename,
                Statistic.TotOfPkg,newPacket->captureLen,(newPacket->BasicInfo).Ip.Src,
                (newPacket->BasicInfo).Port.Src,(newPacket->BasicInfo).Ip.Dst,
                (newPacket->BasicInfo).Port.Dst,(newPacket->BasicInfo).Protocol,Class);
        if(mysql_real_query(&my_connection,query,(unsigned int)strlen(query)))
        {
            printf("Error exec query:%s\n",mysql_error(&my_connection));
            fprintf(PcapLog,"Error exec query:%s\n",mysql_error(&my_connection));
            exit(1);
            fflush(PcapLog);
        }
    }
}
// process a packet in pcap file, record info to struct packet
int process_packet(FILE* fp,struct packet *newPacket,char Class[50])
{
    unsigned int captureLen,ipHeaderLen,ipTotalLen,tcpHeaderLen;
    unsigned long ipSrc,ipDst;
    unsigned short portSrc,portDst;
    int i,MasterID=0,ViceID=0,IPFragment=0;
    skip(fp,8);

    //captureLen
    process_4B_uint(fp, &captureLen);
//    newPacket->captureLen=captureLen;
    skip(fp, 4); //ethernet header
    MasterID=fgetc(fp);
    fgetc(fp);
    ViceID=fgetc(fp);
    skip(fp, 10); //ethernet header
    getClass(MasterID,ViceID,Class);
    if(fgetc(fp)!=IPPROTO_IP||!strcmp(Class,"\"UNKNOWN\""))
    {
        skip(fp, captureLen-14); //packet is second layer protocol,we skip it  
        return 0;
    }
    ipHeaderLen=(unsigned int)(fgetc(fp)& 0x0f)*4;
    //   newPacket->ipHeaderLen=ipHeaderLen;
    if(ipHeaderLen<20)
    {
        skip(fp,captureLen-15);
        return 0;
    }
    // skip(fp, 8);
    fgetc(fp);
    ipTotalLen=16*16*fgetc(fp)+fgetc(fp);
    if(ipTotalLen<ipHeaderLen)
    {
        skip(fp,captureLen-18);    
        return 0;
    }
    //  skip(fp, 5);
    skip(fp,2);
    IPFragment=fgetc(fp)|fgetc(fp);
    fgetc(fp);
    if(IPFragment!=0)
    {
        skip(fp,captureLen-23);
        return 0;
    }
    (newPacket->BasicInfo).Protocol=(unsigned char)fgetc(fp);
    skip(fp,2); //checksum

    process_ip(fp, &ipSrc);
    process_ip(fp, &ipDst);
    (newPacket->BasicInfo).Ip.Src=ipSrc;
    (newPacket->BasicInfo).Ip.Dst=ipDst;
    skip(fp, ipHeaderLen-20);
    if((newPacket->BasicInfo).Protocol==IPPROTO_TCP||(newPacket->BasicInfo).Protocol==IPPROTO_UDP) //it is a tcp or udp packet
    {
        process_port(fp, &portSrc);
        process_port(fp, &portDst);
        (newPacket->BasicInfo).Port.Src=portSrc;
        (newPacket->BasicInfo).Port.Dst=portDst;
        if((newPacket->BasicInfo).Protocol==IPPROTO_TCP)
        {
            skip(fp, 8);
            tcpHeaderLen=(unsigned int)(fgetc(fp)& 0xf0)/4;
            newPacket->captureLen=ipTotalLen-ipHeaderLen-tcpHeaderLen;
            skip(fp, captureLen-27-ipHeaderLen);
        }
        else
        {
            newPacket->captureLen=16*16*fgetc(fp)+fgetc(fp)-8;
            skip(fp, captureLen-20-ipHeaderLen);
        }
        //  skip(fp, captureLen-18-ipHeaderLen);
    }
    else
    {
        skip(fp, captureLen-14-ipHeaderLen); //packet is neither tcp nor udp, we ignore it 
        return 0;
    }
    return 1;
}
/*判断文件是否是个目录文件，是目录返回 1 ，不是目录返回 0*/
int testdir(char *path)
{
    struct stat buf;
    if(lstat(path,&buf)<0)
    {
        return 0;
    }
    if(S_ISDIR(buf.st_mode))
    {
        return 1; //directory
    }
    return 0;
}
int IsPcapFile(FILE* PcapLog,char* filename)
{
    FILE* fp;
    if((fp=fopen(filename,"rb"))==NULL) 
    {
        fclose(fp);
        printf("File %s Open Error...\n",filename);
        fprintf(PcapLog,"File %s Open Error...\n",filename);
        fflush(PcapLog);
        return 0;
    }
    if(fgetc(fp)==212&&fgetc(fp)==195&&fgetc(fp)==178&&fgetc(fp)==161
            &&fgetc(fp)==2&&fgetc(fp)==0
            &&fgetc(fp)==4&&fgetc(fp)==0)
    {
        fclose(fp);
        return 1;
    }
    fclose(fp);
    return 0;
}
/*遍历目录*/
int directory(FILE* PcapLog,char *path)
{
    DIR *db;              /*保存 打开目录类型文件信息的 结构体*/
    char filename[FILE_NAME],file_name[FILE_NAME];        /*文件名*/
    struct stat file_info;         /*保存 普通文件的属性信息的 结构体*/
    //   struct type *head,*current,*temp;
    //char *delim="_",Class[50];
    char Class[50];
    FILE *fp,*fp2;
    struct dirent *dir_info;        /*保存 目录类型文件属性信息的 结构体*/
    struct packet* newPacket; 
    //  struct flow  *FlowHead,*FlowCurrent,*FlowTemp;
    struct statistic Statistic;
    struct timeval start,end;
    char FilenameAndSize[40];
    int i;
    unsigned long timeuse; 
    unsigned long current_tot_time=0;

    memset(&Statistic,0,sizeof(struct statistic));
    newPacket=(struct packet*)malloc(sizeof(struct packet));
    db=opendir(path);
    if(db==NULL){
        printf("Open dir error !\n");
        fprintf(PcapLog,"Open dir error !\n");
        fflush(PcapLog);
        return 0;
    }
    memset(filename,0,FILE_NAME);
    i=0;
    while ((dir_info=readdir(db)))
    {
        if((strcmp(dir_info->d_name,".")==0)||(strcmp(dir_info->d_name,"..")==0))
            continue;
        else
        {
            sprintf(filename,"%s/%s",path,dir_info->d_name);
            printf("%s/%s\n",path,dir_info->d_name);
            fprintf(PcapLog,"%s/%s\n",path,dir_info->d_name);
            if(!testdir(filename)&&IsPcapFile(PcapLog,filename))
            {
                memset( &file_info,0,sizeof(file_info));   /*将文件信息结构体置 0 */
                lstat( filename, &file_info );      /*通过文件路径，把文件信息结构体填满*/
                strcpy(file_name,dir_info->d_name);
                gettimeofday(&start,NULL);
                if((fp=fopen(filename,"rb"))==NULL)
                {
                    printf("> File doesn't exist!\n");
                    fprintf(PcapLog,"> File doesn't exist!\n");
                    fflush(PcapLog);
                    exit(1);
                }
                skip(fp,24);
                sprintf(FilenameAndSize,"%s|Size:%s",
                        file_name,sizeTransform(file_info.st_size));
                i++;
                sprintf(Statistic.filename,"\"%s\"",file_name);
                Statistic.TotOfPkg=0;
                Statistic.TotOfReg=0;
                while(!feof(fp))
                {
                    memset(newPacket,0,sizeof(struct packet));
                    if(process_packet(fp,newPacket,Class))
                    {
                        Statistic.TotOfPkg++;
                        Statistic.TotOfReg++;
                        process_flow(PcapLog,newPacket,Class,Statistic);
                    }
                    else
                        if(!feof(fp))
                            Statistic.TotOfPkg++;
                    printf("%s|Total number of Packages(Reg\\Tot):%d\\%d|Current Protocol is %s       \r",FilenameAndSize,Statistic.TotOfReg,Statistic.TotOfPkg,Class);
                    fflush(stdout);
                }
                fprintf(PcapLog,"%s|Total number of Packages(Reg\\Tot):%d\\%d|Current Protocol is %s",FilenameAndSize,Statistic.TotOfReg,Statistic.TotOfPkg,Class);
                fflush(PcapLog);
                fclose(fp);
                gettimeofday(&end,NULL);
                timeuse=end.tv_sec-start.tv_sec;
                current_tot_time+=timeuse;
                printf("\nThis is No.%d File...And Run time is %s,",i,timeTransform(timeuse));
                printf("Current Total time is %s\n",timeTransform(current_tot_time));
                fprintf(PcapLog,"\nThis is No.%d File...And Run time is %s,",i,timeTransform(timeuse));
                fprintf(PcapLog,"Current Total time is %s\n",timeTransform(current_tot_time));
                fflush(PcapLog);
                fflush(stdout);
            }
        }
        memset(filename,0,64);
    }
    closedir(db);
    return 0;
}

//program starts here
int main(int argc, char *argv[])
{
    struct timeval start,end;
    unsigned long timeuse;
    MYSQL_RES *res;
    MYSQL_ROW row;
    FILE* PcapLog;
    //char cmd[300];
    gettimeofday(&start,NULL);
    PcapLog=fopen("RetrievePcap.log","w");
    if(argc<2)
    {
        printf("pcap_retrieve.exe [The path of directory]\n");
        fprintf(PcapLog,"pcap_retrieve.exe [The path of directory]\n");
        fflush(PcapLog);
        return 0;
    }
    mysql_init(&my_connection);
    if(mysql_real_connect(&my_connection,"localhost","root","111111",
                "database_flow",0,NULL,CLIENT_MULTI_RESULTS))
    {
        if(mysql_query(&my_connection,"drop table if exists table_flow_final"))
        {
            printf("Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fprintf(PcapLog,"Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fflush(PcapLog);
            exit(1);
        }
        if(mysql_query(&my_connection,"create table table_flow_final \
                    (Filename char(30),ID bigint,SrcIP bigint,SrcPort int,DstIP bigint,DstPort int,Protocol int, \
                     PKG_SIZE_1 int,PKG_SIZE_2 int,PKG_SIZE_3 int,PKG_SIZE_4 int,PKG_SIZE_5 int,\
                     FLOW_SIZE bigint,Class char(50),PRIMARY KEY(SrcIP,SrcPort,DstIP,DstPort,Protocol))"))

        {
            printf("Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fprintf(PcapLog,"Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fflush(PcapLog);
            exit(1);
        }

        directory(PcapLog,argv[1]);
        /*        if(mysql_query(&my_connection,"select count(*) from table_flow_final where PKG_SIZE_5>0"))
                  {
                  printf("Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
                  fprintf(PcapLog,"Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
                  fflush(PcapLog);
                  exit(1);
                  }
                  res=mysql_store_result(&my_connection);
                  row=mysql_fetch_row(res);
                  printf("Total available flow number is %s\n",row[0]);
                  fprintf(PcapLog,"Total available flow number is %s\n",row[0]);
                  fflush(PcapLog);
                  mysql_free_result(res);*/
        if(mysql_query(&my_connection,"select count(*) from table_flow_final"))
        {
            printf("Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fprintf(PcapLog,"Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fflush(PcapLog);
            exit(1);
        }
        res=mysql_store_result(&my_connection);
        row=mysql_fetch_row(res);
        printf("Total flow number is %s\n",row[0]);
        fprintf(PcapLog,"Total flow number is %s\n",row[0]);
        fflush(PcapLog);
        mysql_free_result(res);
        if(mysql_query(&my_connection,"select MAX(FLOW_SIZE) from table_flow_final"))
        {
            printf("Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fprintf(PcapLog,"Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fflush(PcapLog);
            exit(1);
        }
        res=mysql_store_result(&my_connection);
        row=mysql_fetch_row(res);
        printf("Max flow size is %s\n",sizeTransform(atoll(row[0])));
        fprintf(PcapLog,"Max flow size is %s\n",sizeTransform(atoll(row[0])));
        fflush(PcapLog);
        mysql_free_result(res);
        if(mysql_query(&my_connection,"select MIN(FLOW_SIZE) from table_flow_final"))
        {
            printf("Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fprintf(PcapLog,"Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fflush(PcapLog);
            exit(1);
        }
        res=mysql_store_result(&my_connection);
        row=mysql_fetch_row(res);
        printf("Min flow size is %s\n",sizeTransform(atoll(row[0])));
        fprintf(PcapLog,"Min flow size is %s\n",sizeTransform(atoll(row[0])));
        fflush(PcapLog);
        mysql_free_result(res);
        if(mysql_query(&my_connection,"select AVG(FLOW_SIZE) from table_flow_final"))
        {
            printf("Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fprintf(PcapLog,"Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fflush(PcapLog);
            exit(1);
        }
        res=mysql_store_result(&my_connection);
        row=mysql_fetch_row(res);
        printf("Average flow size is %s\n",sizeTransform(atoll(row[0])));
        fprintf(PcapLog,"Average flow size is %s\n",sizeTransform(atoll(row[0])));
        fflush(PcapLog);
        mysql_free_result(res);
        if(mysql_query(&my_connection,"select SUM(FLOW_SIZE) from table_flow_final"))
        {
            printf("Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fprintf(PcapLog,"Select error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fflush(PcapLog);
            exit(1);
        }
        res=mysql_store_result(&my_connection);
        row=mysql_fetch_row(res);
        printf("Total Recognized flow size is %s\n",sizeTransform(atoll(row[0])));
        fprintf(PcapLog,"Total Recognized flow size is %s\n",sizeTransform(atoll(row[0])));
        fflush(PcapLog);
        mysql_free_result(res);
        mysql_close(&my_connection);
    }
    else
    {
        printf("Connection failed\n");    
        fprintf(PcapLog,"Connection failed\n");    
        fflush(PcapLog);
        if(mysql_errno(&my_connection))
        {
            printf("Connection error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fprintf(PcapLog,"Connection error %d: %s\n",mysql_errno(&my_connection),mysql_error(&my_connection));
            fflush(PcapLog);
        }
    }
    gettimeofday(&end,NULL);
    timeuse=(end.tv_sec-start.tv_sec);
    printf("Run time is %s\n",timeTransform(timeuse));
    fprintf(PcapLog,"Run time is %s\n",timeTransform(timeuse));
    fflush(PcapLog);
    fclose(PcapLog);
    return 0;
}
